import os

import keras.losses
from keras import optimizers
from keras.models import Model

from custom_layer import MyDense


class MyModel(Model):
    """
    自定义网络模型
    """

    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = MyDense(28 * 28, 256)
        self.fc2 = MyDense(256, 128)
        self.fc3 = MyDense(128, 32)
        self.fc4 = MyDense(32, 10)

    def call(self, inputs, training=None, mask=None):
        fc1_output = self.fc1.call(inputs=inputs)
        fc2_output = self.fc2.call(inputs=fc1_output)
        fc3_output = self.fc3.call(inputs=fc2_output)
        fc4_output = self.fc4.call(inputs=fc3_output)

        return fc4_output


if __name__ == '__main__':
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

    model = MyModel()

    model.compile(optimizer=optimizers.Adam(learning_rate=0.01),
                  loss=keras.losses.CategoricalCrossentropy(from_logits=True), metrics=['accuracy'])

    model.build(input_shape=[None, 28 * 28])
    model.summary()
